#include <iostream>
#include <sstream>

#include <glog/logging.h>

#include <QtCore>

#include <botan/rsa.h>
#include <botan/symkey.h>
#include <botan/pubkey.h>
#include <botan/zlib.h>

#include "octopod.h"
#include "crypto.h"

namespace octopod {

// Currently doesn't handle different Botan errors differently.
Botan::SymmetricKey *crypto::symmetricKeyGen() {
  try {
    Botan::AutoSeeded_RNG rng;
    return new Botan::SymmetricKey(rng, crypto::kKeyLength);
  } catch (Botan::Exception &e) {
    DLOG(ERROR) << Q_FUNC_INFO << "Botan exception: " << e.what();
  }

  return NULL;
}

// Doesn't handle different Botan errors differently.
Botan::Private_Key *crypto::asymmetricPrivateKeyGen() {
  try {
    Botan::AutoSeeded_RNG rng;
    return new Botan::RSA_PrivateKey(rng, crypto::kRSABits);
  } catch (Botan::Exception &e) {
    DLOG(ERROR) << Q_FUNC_INFO << "Botan exception: " << e.what();
  }

  return NULL;
}

// Generates a public key by casting the private key, as per Botan docs.
const Botan::Public_Key *crypto::getPublicKey(const Botan::Private_Key *private_key) {
  CHECK_NOTNULL(private_key);
  return dynamic_cast<const Botan::Public_Key *>(private_key);
}

// Creates a temporary buffer to store the random bytes.
// May be inefficient for large buffers.
bool crypto::randomGen(int num, std::string *s) {
  CHECK(num > 0);
  CHECK_NOTNULL(s);

  try {
    Botan::AutoSeeded_RNG rng;
    Botan::byte buf[num];

    rng.randomize(buf, num);
    s->assign(reinterpret_cast<const char *>(buf), num);
    return true;
  } catch (Botan::Exception &e) {
    DLOG(ERROR) << Q_FUNC_INFO << "Botan exception: " << e.what();
  }

  return false;
}

// Does a one-shot authenticated encryption.
// Uses a temporary string, so it's not that efficient.
bool crypto::authEncrypt(const Botan::SymmetricKey &key, const std::string &plaintext, std::string *ciphertext, std::string *init_vector) {
  CHECK_NOTNULL(ciphertext);
  CHECK_NOTNULL(init_vector);

  try {
    Botan::AutoSeeded_RNG rng;
    Botan::InitializationVector iv(rng, kIVLength);

    // store the IV
    init_vector->assign(reinterpret_cast<const char *>(iv.begin()), kIVLength);
  
    // generate and store the ciphertext
    std::ostringstream out(std::ios::binary | std::ios::out);
    Botan::Pipe pipe(new Botan::Zlib_Compression(3), get_cipher(kAECipher, key, iv, Botan::ENCRYPTION), new Botan::DataSink_Stream(out));
    pipe.process_msg(plaintext);
    *ciphertext = out.str();

    return true;
  } catch (Botan::Exception &e) {
    DLOG(ERROR) << Q_FUNC_INFO << "Botan exception: " << e.what();
  }

  return false;
}

// Does a one-shot authenticated decryption.
// Uses a temporary string, so it's not that efficient.
// By virtue of Botan throwing an exception if authentication fails,
// this will return false if anything during decryption is amiss.
bool crypto::authDecrypt(const Botan::SymmetricKey &key, const std::string &ciphertext, const std::string &init_vector, std::string *plaintext) {
  CHECK_NOTNULL(plaintext);

  try {
    Botan::InitializationVector iv(reinterpret_cast<const Botan::byte *>(init_vector.data()), init_vector.size());

    // decrypt
    std::ostringstream out(std::ios::binary | std::ios::out);
    Botan::Pipe pipe(get_cipher(kAECipher, key, iv, Botan::DECRYPTION), new Botan::Zlib_Decompression, new Botan::DataSink_Stream(out));
    pipe.process_msg(ciphertext);
    *plaintext = out.str();

    return true;
  } catch (Botan::Exception &e) {
    DLOG(ERROR) << Q_FUNC_INFO << "Botan exception: " << e.what();
  }

  return false;
}

// Requires that the underlying public key is large enough to encrypt things
// of the size of the given plaintext.
bool crypto::pubkeyEncrypt(const Botan::Public_Key *public_key, const std::string &plaintext, std::string *ciphertext) {
  try {
    Botan::AutoSeeded_RNG rng;
    Botan::PK_Encryptor_EME enc(*public_key, std::string(kPublicKeyEncryptionMethod));

    CHECK(plaintext.size() <= enc.maximum_input_size());

    Botan::SecureVector<Botan::byte> ctext = enc.encrypt(reinterpret_cast<const Botan::byte *>(plaintext.data()), plaintext.size(), rng);
    ciphertext->assign(reinterpret_cast<const char *>(ctext.begin()), ctext.size());

    return true;
  } catch (Botan::Exception &e) {
    DLOG(ERROR) << Q_FUNC_INFO << "Botan exception: " << e.what();
  }

  return false;
}

// Fills the string plaintext with the output of a one-shot
// public key decryption call using the given private key and ciphertext.
// Returns true if the decryption succeeds.
bool crypto::pubkeyDecrypt(const Botan::Private_Key *private_key, const std::string &ciphertext, std::string *plaintext) {
  try {
    Botan::AutoSeeded_RNG rng;
    Botan::PK_Decryptor_EME dec(*private_key, kPublicKeyEncryptionMethod);

    Botan::SecureVector<Botan::byte> ptext = dec.decrypt(reinterpret_cast<const Botan::byte *>(ciphertext.data()), ciphertext.size());
    plaintext->assign(reinterpret_cast<const char *>(ptext.begin()), ptext.size());

    return true;
  } catch (Botan::Exception &e) {
    DLOG(ERROR) << Q_FUNC_INFO << "Botan exception: " << e.what();
  }

  return false;
}

// Should never fail to produce a hash result.
std::string crypto::hash(const std::string &data) {
  try {
    Botan::Pipe pipe(new Botan::Hash_Filter(kHash));
    pipe.process_msg(data);
    
    return pipe.read_all_as_string();
  } catch (Botan::Exception &e) {
    DLOG(ERROR) << Q_FUNC_INFO << "Botan exception: " << e.what();
  }

  LOG(FATAL) << Q_FUNC_INFO << "Failed to hash input data of length " << data.size();
  return std::string();
}

bool crypto::MAC(const Botan::SymmetricKey &key, const std::string &data, std::string *tag) {
  CHECK_NOTNULL(tag);

  try {
    Botan::Pipe pipe(new Botan::MAC_Filter(kMAC, key));
    pipe.process_msg(data);

    *tag = pipe.read_all_as_string();

    return true;
  } catch (Botan::Exception &e) {
    DLOG(ERROR) << Q_FUNC_INFO << "Botan exception: " << e.what();
  }

  return false;
}

bool crypto::KDF(const Botan::SymmetricKey &key, const std::string &salt, Botan::SymmetricKey *derived_key) {
  CHECK_NOTNULL(derived_key);

  std::string tag;
  if (!MAC(key, salt, &tag)) {
    DLOG(WARNING) << Q_FUNC_INFO << "key derivation failed.";
    return false;
  }

  *derived_key = Botan::SymmetricKey(reinterpret_cast<const Botan::byte *>(tag.data()), tag.size());

  return true;
}

}  // namespace octopod
